/**
 * @file pcap_krb5.c
 * @brief
 * @author thewangcj
 * @version 0.1
 * @date 2022-02-09 18:02:21
 */
#include "pcap_krb5.h"
#include "krb5_asn1_tab.h"
#include <arpa/inet.h>
#include <libtasn1.h>
#include <malloc.h>
#include <memory.h>
#include <netinet/in.h>
#include <pcap.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#define DEBUG
#pragma pack(1)

#define UP_OR_DOWN(a, b) ((a->s_port == b->s_port) ? 1 : 0)

static void _extract_kerberos_errcode(asn1_node *element)
{
    int ret = 0;
    char data[8];
    memset(data, 0, 8);
    int len = sizeof(data);
    ret     = asn1_read_value(*element, "error-code", data, &len);
    if (ret != ASN1_SUCCESS) {
        return;
    } else {
        char err_code[8];
        sprintf(err_code, "%d", data[0]);
    }
}

static void _extract_kerberos_cname(asn1_node *element, uint64_t tag)
{
    int ret = 0;
    char data[1024];
    memset(data, 0, 1024);
    int len     = sizeof(data);
    int ele_num = 0;
    int i       = 0;
    char name[MAX_NAME_LEN];
    switch (tag) {
    case KERBEROS_AS_REQ:
    case KERBEROS_TGS_REQ:
        // 获取 cname 的数量
        ret = asn1_number_of_elements(*element, "req-body.cname.name-string", &ele_num);
        if (ret != ASN1_SUCCESS) {
            break;
        }
        // 遍历获取 cname
        for (i = 0; i < ele_num; i++) {
            sprintf(name, "req-body.cname.name-string.?%d", i + 1);
            ret = asn1_read_value(*element, name, data, &len);
            if (ret != ASN1_SUCCESS) {
                break;
            } else {
                printf("cname : %s\n", data);
            }
        }
        break;
    case KERBEROS_AS_REP:
    case KERBEROS_TGS_REP:
        ret = asn1_number_of_elements(*element, "cname.name-string", &ele_num);
        if (ret != ASN1_SUCCESS) {
            break;
        }
        for (i = 0; i < ele_num; i++) {
            sprintf(name, "cname.name-string.?%d", i + 1);
            ret = asn1_read_value(*element, name, data, &len);
            if (ret != ASN1_SUCCESS) {
                break;
            } else {
                printf("cname : %s\n", data);
            }
        }
        break;
    case KERBEROS_ERROR: break;
    default: break;
    }
}

static void _extract_kerberos_realm(asn1_node *element, uint64_t tag)
{
    int ret = 0;
    char data[1024];
    memset(data, 0, 1024);
    int len = sizeof(data);
    switch (tag) {
    case KERBEROS_AS_REQ:
    case KERBEROS_TGS_REQ: ret = asn1_read_value(*element, "req-body.realm", data, &len); break;
    case KERBEROS_AS_REP:
    case KERBEROS_TGS_REP:
    case KERBEROS_AP_REQ: ret = asn1_read_value(*element, "ticket.realm", data, &len); break;
    case KERBEROS_ERROR: ret = asn1_read_value(*element, "realm", data, &len); break;
    default: break;
    }
    if (ret != ASN1_SUCCESS) {
        return;
    } else {
        printf("realm : %s\n", data);
    }
}

static void _extract_kerberos_sname(asn1_node *element, uint64_t tag)
{
    int ret = 0;
    char data[1024];
    memset(data, 0, 1024);
    int len     = sizeof(data);
    int ele_num = 0;
    int i       = 0;
    char name[MAX_NAME_LEN];
    switch (tag) {
    case KERBEROS_AS_REQ:
    case KERBEROS_TGS_REQ:
        ret = asn1_number_of_elements(*element, "req-body.sname.name-string", &ele_num);
        if (ret != ASN1_SUCCESS) {
            return;
        }
        // 遍历获取 sname
        for (i = 0; i < ele_num; i++) {
            sprintf(name, "req-body.sname.name-string.?%d", i + 1);
            ret = asn1_read_value(*element, name, data, &len);
            if (ret != ASN1_SUCCESS) {
                break;
            } else {
                printf("sname : %s\n", data);
            }
        }
        break;
    case KERBEROS_AS_REP:
    case KERBEROS_TGS_REP:
    case KERBEROS_AP_REQ:
        ret = asn1_number_of_elements(*element, "ticket.sname.name-string", &ele_num);
        if (ret != ASN1_SUCCESS) {
            return;
        }
        // 遍历获取 sname
        for (i = 0; i < ele_num; i++) {
            sprintf(name, "ticket.sname.name-string.?%d", i + 1);
            ret = asn1_read_value(*element, name, data, &len);
            if (ret != ASN1_SUCCESS) {
                break;
            } else {
                printf("sname : %s\n", data);
            }
        }
        break;
    case KERBEROS_ERROR:
        ret = asn1_number_of_elements(*element, "sname.name-string", &ele_num);
        if (ret != ASN1_SUCCESS) {
            return;
        }
        // 遍历获取 sname
        for (i = 0; i < ele_num; i++) {
            sprintf(name, "sname.name-string.?%d", i + 1);
            ret = asn1_read_value(*element, name, data, &len);
            if (ret != ASN1_SUCCESS) {
                break;
            } else {
                printf("sname : %s\n", data);
            }
        }
        break;
    default: break;
    }
}

static long _extract_kerberos_msg_type(PPCTX pctx)
{
    unsigned char cls;
    uint64_t tag;
    int ret = 0;
    int len = 0;
    // 跳过 Record Mark
    char *der   = pctx->payload;
    int der_len = pctx->payload_len;
    if (pctx->payload[0] == 0 && pctx->payload[1] == 0) {
        der += 4;
        der_len -= 4;
    }
    ret = asn1_get_tag_der(der, der_len, &cls, &len, &tag);
    if (ret == ASN1_SUCCESS) {
        char tag_str[8];
        sprintf(tag_str, "%ld", tag);
        printf("tag : %s\n", tag_str);
    }
    return tag;
}

static void _extract_ap_req(asn1_node *asn1_element, asn1_node *definitions, char *error_description)
{
    char data[2048];
    int ret = 0;
    memset(data, 0, 2048);
    int len = sizeof(data);
    ret     = asn1_read_value(*asn1_element, "padata.?1.padata-value", data, &len);
    if (ret != ASN1_SUCCESS) {
        return;
    } else {
        asn1_node ap_req = NULL;
        ret              = asn1_create_element(*definitions, "KerberosV5Spec2.AP-REQ", &ap_req);
        if (ret != ASN1_SUCCESS)
            return;
        else {
            ret = asn1_der_decoding2(&ap_req, data, &len, ASN1_DECODE_FLAG_ALLOW_PADDING | ASN1_DECODE_FLAG_STRICT_DER,
                                     error_description);
            if (ret != ASN1_SUCCESS)
                return;
            else {
                uint64_t tag = KERBEROS_AP_REQ;
                char tag_str[8];
                sprintf(tag_str, "%ld", tag);
                printf("tag : %s\n", tag_str);
                _extract_kerberos_realm(&ap_req, tag);
                _extract_kerberos_sname(&ap_req, tag);
                ret = asn1_delete_structure(&ap_req);
                if (ret != ASN1_SUCCESS) {
                    printf("asn1_delete_structure ap_req fail\n");
                }
            }
        }
    }
}

static int _extract_kerberos(PPCTX pctx)
{
    printf("-----------\n");
    int ret                = 0;
    asn1_node asn1_element = NULL;
    asn1_node definitions  = NULL;
    char error_description[ASN1_MAX_ERROR_DESCRIPTION_SIZE];
    // libtasn14.10 版本 asn1_array2tree 线程不安全
    ret = asn1_array2tree(krb5_asn1_tab, &definitions, error_description);
    if (ret != ASN1_SUCCESS) return -1;

    uint64_t tag = _extract_kerberos_msg_type(pctx);
    switch (tag) {
    case KERBEROS_AS_REQ: ret = asn1_create_element(definitions, "KerberosV5Spec2.AS-REQ", &asn1_element); break;
    case KERBEROS_AS_REP: ret = asn1_create_element(definitions, "KerberosV5Spec2.AS-REP", &asn1_element); break;
    case KERBEROS_TGS_REQ: ret = asn1_create_element(definitions, "KerberosV5Spec2.TGS-REQ", &asn1_element); break;
    case KERBEROS_TGS_REP: ret = asn1_create_element(definitions, "KerberosV5Spec2.TGS-REP", &asn1_element); break;
    case KERBEROS_ERROR: ret = asn1_create_element(definitions, "KerberosV5Spec2.KRB-ERROR", &asn1_element);
    default: break;
    }

    if (ret != ASN1_SUCCESS) return -1;

    char buffer[10 * 1024];
    int mark_len = 0;
    if (pctx->payload[0] == 0 && pctx->payload[1] == 0) mark_len = 4;
    int size = pctx->payload_len - mark_len;
    memcpy(buffer, pctx->payload + mark_len, pctx->payload_len - mark_len);
    ret = asn1_der_decoding2(&asn1_element, buffer, &size, ASN1_DECODE_FLAG_ALLOW_PADDING | ASN1_DECODE_FLAG_STRICT_DER,
                             error_description);
    if (ret != ASN1_SUCCESS) {
        // asn1_create_element 后不 delete 会造成内存泄漏
        ret = asn1_delete_structure(&definitions);
        if (ret != ASN1_SUCCESS) {
            printf("asn1_delete_structure definitions fail\n");
        }
        return -1;
    }

    switch (tag) {
    case KERBEROS_AS_REQ:
    case KERBEROS_AS_REP:
    case KERBEROS_TGS_REP:
        _extract_kerberos_cname(&asn1_element, tag);
        _extract_kerberos_realm(&asn1_element, tag);
        _extract_kerberos_sname(&asn1_element, tag);
    case KERBEROS_TGS_REQ:
        // TGS_REQ 中可能还含有 AP_REQ
        _extract_ap_req(&asn1_element, &definitions, error_description);
        break;
    case KERBEROS_ERROR:
        _extract_kerberos_errcode(&asn1_element);
        _extract_kerberos_realm(&asn1_element, tag);
        _extract_kerberos_sname(&asn1_element, tag);
        break;
    default: break;
    }
    printf("-----------\n");
    ret = asn1_delete_structure(&definitions);
    if (ret != ASN1_SUCCESS) {
        printf("asn1_delete_structure definitions fail\n");
    }
    ret = asn1_delete_structure(&asn1_element);
    if (ret != ASN1_SUCCESS) {
        printf("asn1_delete_structure asn1_element fail\n");
    }
    return ret;
}

void process_ip(PPCTX pctx)
{
    const struct packet_ip *ip; /* The IP header */

    /* IP头 */
    ip = (struct packet_ip *)(pctx->packet + SIZE_ETHERNET);
    u_int size_ip;
    size_ip = IP_HL(ip) * 4;
    if (size_ip < 20) {
        printf("无效的IP头长度: %u bytes\n", size_ip);
        return;
    }
    pctx->ip_head_len = size_ip;
    switch (ip->ip_p) {
    case IPPROTO_TCP: process_tcp(pctx); break;
    case IPPROTO_UDP: process_udp(pctx); break;
    default: break;
    }
}

void process_tcp(PPCTX pctx)
{
#ifdef DEBUG
    printf("process_tcp\n");
#endif
    /* TCP头 */
    const struct packet_tcp *tcp; /* The TCP header */
    tcp = (struct packet_tcp *)(pctx->packet + SIZE_ETHERNET + pctx->ip_head_len);
    u_int size_tcp;
    size_tcp = TH_OFF(tcp) * 4;
    if (size_tcp < 20) {
        printf("无效的TCP头长度: %u bytes\n", size_tcp);
        return;
    } else {
        printf("TCP头长度: %u bytes\n", size_tcp);
    }
}

void process_udp(PPCTX pctx)
{
    /* UDP头 */
    const struct packet_udp *udp; /* The UDP header */
    udp = (struct packet_udp *)(pctx->packet + SIZE_ETHERNET + pctx->ip_head_len);
    // uh_len = UDP 头部长度 + 数据部分长度
    u_int data_len;
    data_len = ntohs(udp->uh_len) - UDP_HEADER_LEN;
    printf("krb5 数据长度: %u bytes\n", data_len);
    pctx->payload_len = data_len;
    pctx->payload = pctx->packet + SIZE_ETHERNET + pctx->ip_head_len + UDP_HEADER_LEN;
    _extract_kerberos(pctx);
}

void loop_callback(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{
    /* 以太网头 */
    const struct packet_ethernet *ethernet; /* The ethernet header */

    ethernet = (struct packet_ethernet *)(packet);
    struct pctx p;
    p.packet = packet;
    p.len    = header->len;
    switch (ethernet->ether_type) {
        // 只处理 IPv4
    case IPV4: process_ip(&p); break;
    default: break;
    }
}

void parse_args(int argc, char const *argv[], char **file_name)
{
    int32_t opt = 0;
    while ((opt = getopt(argc, (char *const *)argv, "hvf:i:o:")) != -1) {
        switch (opt) {
        case 'h':
            printf("使用方法：pcap_krb5 [option]:\n");
            printf("-f, 读取 pcap 文件并并解析其中的 KRB5 信息，如：pcap_krb5 -f pcap文件名\n");
            printf("-h, 输出使用方法\n");
            printf("-v, 查看版本信息\n");
            break;
        case 'v': printf("pcap_krb5 v0.1.0\n"); break;
        case 'f':
            *file_name = optarg;
            printf("file_name is %s \n", optarg);
            break;
        default: printf("请输入pcap_krb5 -h 查看使用方法\n"); break;
        }
    }
}

int main(int argc, char const *argv[])
{
    char *file_name;
    parse_args(argc, argv, &file_name);
    if (file_name == NULL) {
        return 0;
    }

    clock_t clock_start = clock();
    char err_buf[PCAP_ERRBUF_SIZE];
    pcap_t *handle;
    handle = pcap_open_offline(file_name, err_buf);
    if (handle == NULL) {
        fprintf(stderr, "无法打开: %s， 请输入正确的 pcap 文件名\n", file_name);
        return -1;
    }
    pcap_loop(handle, -1, loop_callback, NULL);

    pcap_close(handle);

    clock_t clock_end = clock();
    double spend_time = 1000.0 * (clock_end - clock_start) / CLOCKS_PER_SEC;
    printf("解析共花费 %.2lf 毫秒\n", spend_time);
    return 0;
}
